{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Mandelbrot.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "swift",
      "display_name": "Swift"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u5L6EZAoRvji",
        "colab_type": "text"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors. [Licensed under the Apache License, Version 2.0](#scrollTo=Afd8bu4xJOgh)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2o3jcY7CSGkX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "// #@title Licensed under the Apache License, Version 2.0 (the \"License\"); { display-mode: \"form\" }\n",
        "// Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "// you may not use this file except in compliance with the License.\n",
        "// You may obtain a copy of the License at\n",
        "//\n",
        "// https://www.apache.org/licenses/LICENSE-2.0\n",
        "//\n",
        "// Unless required by applicable law or agreed to in writing, software\n",
        "// distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "// See the License for the specific language governing permissions and\n",
        "// limitations under the License."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mNGn0Bj2SNO6",
        "colab_type": "text"
      },
      "source": [
        "# The Mandelbrot set\n",
        "\n",
        "This notebook provides an example of how to calculate and visualize the Mandelbrot set using Swift and the accelerated Tensor type."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "53L3CYBuSwWo",
        "colab_type": "text"
      },
      "source": [
        "## Attaching to the CPU or an accelerator"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kZRlD4utdPuX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import Foundation\n",
        "import TensorFlow"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BDi-9zvIvlfi",
        "colab_type": "text"
      },
      "source": [
        "Let's verify what accelerator, if any, this instance is using. We'll set up a device using the XLA-based X10 backend. ([This tutorial](https://colab.research.google.com/github/tensorflow/swift/blob/main/docs/site/tutorials/introducing_x10.ipynb) has more information on the two backends.)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ebc_D8qrqmvD",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "let device = Device.defaultXLA\n",
        "// let device = Device.defaultTFEager\n",
        "print(\"\\u{001B}[2J\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V-yjuBsrTybf",
        "colab_type": "text"
      },
      "source": [
        "Then, we'll query that device to see what accelerator it found. If none, it will use the CPU."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vjWRZ6ZEquL5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "device"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PxoTq9JDvy6X",
        "colab_type": "text"
      },
      "source": [
        "## Configuring complex number tensors\n",
        "\n",
        "The Mandelbrot set is calculated using complex numbers. To represent a matrix of complex numbers, we'll create a ComplexTensor type that contains separate Tensors for the real and imaginary components.\n",
        "\n",
        "A precondition that the `real` and `imaginary` tensor shapes match would make sense in practice, but for simplicity we won't verify that here."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HBI_jvtgt10o",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "struct ComplexTensor {\n",
        "  let real: Tensor<Float>\n",
        "  let imaginary: Tensor<Float>\n",
        "}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gx3eOtxNyXTC",
        "colab_type": "text"
      },
      "source": [
        "We can then define the mathematical operators that we'll need to work over this new parallel complex type."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8P5Wu5Yyyqbp",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "func +(lhs: ComplexTensor, rhs: ComplexTensor) -> ComplexTensor {\n",
        "  let real = lhs.real + rhs.real\n",
        "  let imaginary = lhs.imaginary + rhs.imaginary\n",
        "  return ComplexTensor(real: real, imaginary: imaginary)\n",
        "}\n",
        "\n",
        "func *(lhs: ComplexTensor, rhs: ComplexTensor) -> ComplexTensor {\n",
        "  let real = lhs.real .* rhs.real - lhs.imaginary .* rhs.imaginary\n",
        "  let imaginary = lhs.real .* rhs.imaginary + lhs.imaginary .* rhs.real\n",
        "  return ComplexTensor(real: real, imaginary: imaginary)\n",
        "}\n",
        "\n",
        "func abs(_ value: ComplexTensor) -> Tensor<Float> {\n",
        "  return value.real .* value.real + value.imaginary .* value.imaginary\n",
        "}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n0j6YY82ymmr",
        "colab_type": "text"
      },
      "source": [
        "And then test them out:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lJ7IDzpMzCQr",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "let complex1 = ComplexTensor(\n",
        "  real: Tensor<Float>([1.0, 2.0], on: device),\n",
        "  imaginary: Tensor<Float>([2.0, 3.0], on: device))\n",
        "let complex2 = ComplexTensor(\n",
        "  real: Tensor<Float>([2.5, 3.5], on: device),\n",
        "  imaginary: Tensor<Float>([3.5, 4.5], on: device))\n",
        "complex1 + complex2"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RcijtTxvzP9A",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "complex1 * complex2"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "h3uFopdHzmx0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "abs(complex1)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C_hWCnvPy2EH",
        "colab_type": "text"
      },
      "source": [
        "## Calculating the Mandelbrot set\n",
        "\n",
        "We can set up the parameters for the space of complex numbers to calculate over:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A3ivlmU7vPNd",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "let xDensity = 1030\n",
        "let yDensity = 1030\n",
        "let tolerance: Float = 4.0\n",
        "let realMinimum: Float = -2.0\n",
        "let realMaximum: Float = 1.0\n",
        "let imaginaryMinimum: Float = -1.3\n",
        "let imaginaryMaximum: Float = 1.3\n",
        "let iterations: Int32 = 200"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0ev1ruhf02sP",
        "colab_type": "text"
      },
      "source": [
        "From the above parameters, ComplexTensors can be initialized that span the complex number space. We also need to configure the variables used in the calculation."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ydjUMfZq0e_Z",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "let xs = Tensor<Float>(linearSpaceFrom: realMinimum, to: realMaximum, count: xDensity, on: device)\n",
        "  .broadcasted(to: [xDensity, yDensity])\n",
        "let ys = Tensor<Float>(linearSpaceFrom: imaginaryMinimum, to: imaginaryMaximum, count: yDensity, on: device)\n",
        "  .expandingShape(at: 1).broadcasted(to: [xDensity, yDensity])\n",
        "let X = ComplexTensor(real: xs, imaginary: ys)\n",
        "var Z = ComplexTensor(real: Tensor(zerosLike: xs), imaginary: Tensor(zerosLike: ys))\n",
        "var divergence = Tensor<Float>(repeating: Float(iterations), shape: xs.shape, on: device)\n",
        "\n",
        "// For X10, we'll make sure the initialization of these tensors doesn't carry\n",
        "// into the trace for the first iteration.\n",
        "LazyTensorBarrier()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ogyf62f-zLBS",
        "colab_type": "text"
      },
      "source": [
        "Finally, we can perform the actual Mandelbrot set calculation for a specified number of iterations. We mark the iteration at which the calculation diverges, if it does, for later visualization."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SDlnlp7GwiRm",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "let start = Date()\n",
        "\n",
        "for iteration in 0..<iterations {\n",
        "  Z = Z * Z + X\n",
        "\n",
        "  let aboveThreshold = abs(Z) .> tolerance\n",
        "  divergence = divergence.replacing(with: min(divergence, Float(iteration)), where: aboveThreshold)\n",
        "\n",
        "  // For X10, we're cutting the trace to be a single iteration.\n",
        "  LazyTensorBarrier()\n",
        "}\n",
        "\n",
        "print(\"Total calculation time: \\(String(format: \"%.3f\", Date().timeIntervalSince(start))) seconds\")\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lDpv-ljz2NSS",
        "colab_type": "text"
      },
      "source": [
        "## Visualizing the Mandelbrot set\n",
        "\n",
        "At each point in the complex space we operated over, `divergence` now contains the iteration at which the calculation diverged for that starting complex number. If it did not, it contains the maximum number of iterations.\n",
        "\n",
        "We'll pull in Matplotlib to help us visualize the results inline:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NBfe0Jzvw7m1",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import PythonKit\n",
        "%include \"EnableIPythonDisplay.swift\"\n",
        "IPythonDisplay.shell.enable_matplotlib(\"inline\")\n",
        "let plt = Python.import(\"matplotlib.pyplot\")\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QAcVjfqf0XIt",
        "colab_type": "text"
      },
      "source": [
        "set up a function to plot the set over a given area:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "J4zJMEQW0bEc",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "func plotMandelbrot(_ divergence: Tensor<Float>, xLimits: (Float, Float), yLimits: (Float, Float)) {\n",
        "  let fig = plt.figure(figsize: [10, 10])\n",
        "\n",
        "  let ax = fig.add_axes([0.0, 0.0, 1.0, 1.0], frameon: false, aspect: 1)\n",
        "  ax.set_xticks(Array<Int>([]))\n",
        "  ax.set_yticks(Array<Int>([]))\n",
        "  ax.set_xlim(xLimits.0, xLimits.1)\n",
        "  ax.set_ylim(yLimits.0, yLimits.1)\n",
        "\n",
        "  let colorMap = plt.get_cmap(\"prism\")\n",
        "  colorMap.set_over(\"black\")\n",
        "\n",
        "  plt.imshow(divergence.scalars.makeNumpyArray().reshape([xDensity, yDensity]),\n",
        "    extent: [realMinimum, realMaximum, imaginaryMinimum, imaginaryMaximum],\n",
        "    cmap: colorMap, vmax: (iterations - 1))\n",
        "\n",
        "  plt.show()  \n",
        "}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JVgMGx5P2YLP",
        "colab_type": "text"
      },
      "source": [
        "and then display the calculated Mandelbrot set at a few magnifications:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_-5hX3mR2Hnv",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "plotMandelbrot(divergence, xLimits: (realMinimum, realMaximum), yLimits: (imaginaryMinimum, imaginaryMaximum))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sfmE4MkKv-Bx",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "plotMandelbrot(divergence, xLimits: (-1.0, 0.0), yLimits: (0.0, 1.0))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jfBNV4vOyS4E",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "plotMandelbrot(divergence, xLimits: (-1.0, -0.5), yLimits: (0.0, 0.4))"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}
